import time
import socket
import requests
from qkdKey import QKDKey
from qkdBuffer import QKDBuffer

class QKDKeyGen:
    def __init__(self, key_size, key_rate, packet_size, data_rate, max_sifting_packets, master_address, slave_address):
        self.key_size = key_size  # Generated key size
        self.key_rate = key_rate  # Number of times to execute generate_and_distribute_keys per second
        self.packet_size = packet_size  # Network packet size
        self.data_rate = data_rate  # Network traffic intensity
        self.max_sifting_packets = max_sifting_packets  # Max sifting packets size
        self.master_address = master_address  # Local KMS address (Alice)
        self.slave_address = slave_address  # Remote KMS address (Bob)
        self.counter = 0

    def start_communication(self):
        # Generate and add keys to both KMS systems based on key rate
        start_time = time.time()
        while True:
            self.generate_and_distribute_keys()
            self.counter += 1
            # Calculate time taken and adjust to maintain the key rate per second
            elapsed_time = time.time() - start_time
            if self.counter >= self.key_rate:
                sleep_time = max(1 - elapsed_time, 0)
                time.sleep(sleep_time)
                start_time = time.time()
                self.counter = 0

    def generate_and_distribute_keys(self):
        key = QKDKey.generate_random_key(self.key_size)
        self.send_key_to_kms(self.master_address, key)
        self.send_key_to_kms(self.slave_address, key)

    def send_key_to_kms(self, kms_address, key):
        try:
            url = f"http://{kms_address}:8080/api/keys/add_key"
            print(f"tag1 url:{url}")
            data = {
                "ip_address": kms_address,
                "key_value": key.key_value.hex(),
                "key_id": key.key_id
            }
            response = requests.post(url, json=data)
            if response.status_code == 200:
                print(f"Key added to KMS at {kms_address}. Key ID: {key.key_id}")
            else:
                print(f"Failed to add key to KMS at {kms_address}. Status Code: {response.status_code}")
        except requests.exceptions.RequestException as e:
            print(f"Error sending key to KMS at {kms_address}: {str(e)}")

# if __name__ == "__main__":
#     key_gen = QKDKeyGen(key_size=256, key_rate=5, packet_size=1024, data_rate=100, max_sifting_packets=10, master_address="192.168.1.1", slave_address="192.168.1.2")
#     key_gen.start_communication()
